/*
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */
package net.javlov.vf.fa;

import java.io.Serializable;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.List;

import net.javlov.Option;
import net.javlov.State;
import net.javlov.vf.IneligibleOptionException;
import net.javlov.vf.OptionValues;
import net.javlov.vf.QFunction;
import net.javlov.vf.Value;

/**
 * Abstract superclass for a linear Q-function approximator. The abstract superclass stores the weights and available
 * options; how the weights are updated is left to the subclass.
 * 
 * This class expects the received states to be already encoded into features; this encoding therefore needs to happen
 * externally. The advantage of this is that the encoding, if required by multiple classes, need happen only once.
 * 
 * @author Matthijs Snel
 *
 */
public abstract class QFA<T> implements QFunction<T>{

	private static final long serialVersionUID = 4311740613242908912L;
	
	/**
	 * Available options, regardless of state.
	 */
	protected List<? extends Option<?>> optionPool;
	
	/**
     * Function weights. One row (1st dim) per option, one column (2nd dim) per feature.
     */
    protected double[][] weights;
    
    /**
     * Eligibility traces.
     */
    protected LinearFATraces traces;
    
    /**
     * Dimensionality of the feature vector.
     */
    protected int numFeats;
    
    QFA(int numFeats, List<? extends Option<?>> optionPool) {
    	this.optionPool = optionPool;
    	this.numFeats = numFeats;
    	init();    	
    }
    
	@Override
	public double getValue(State<? extends T> s, Option<?> o) {
		try {
			return getValues(s).get(o);
		} catch ( IneligibleOptionException e ) {
			throw new IneligibleOptionException(o + " not eligible in " + s);
		}
	}

	@Override
	public void init() {
		weights = new double[optionPool.size()][numFeats];
	}

	@Override
	public Double queryValue(State<? extends T> s, Option<?> o) {
		return getValue(s,o);
	}

	@Override
	public void reset() {
		//do nothing		
	}

	@Override
	public void setValue(State<? extends T> s, Option<?> o, double value) {
		throw new UnsupportedOperationException();		
	}

	@Override
	public void update(State<? extends T> s, Option<?> o, double delta) {
		getValues(s).update(o, delta);
	}
	
	public void setTraces(LinearFATraces traces) {
		this.traces = traces;
	}
	
	public LinearFATraces getTraces() {
		return traces;
	}
	
	/**
	 * Abstract superclass for function-approximation-based option values. This class takes care of the
	 * {@link #getMaxOptions()}, {@link #getValues()} and {@link #max()} operations, while the
	 * {@link #get(Option)} and {@link #update(Option, double)} operations should be implemented by
	 * the subclass.
	 * 
	 * @author Matthijs Snel
	 *
	 */
	public abstract class FAOptionValues implements Serializable, OptionValues {

    	private static final long serialVersionUID = -175768075273183858L;
    	
    	/**
    	 * The encoding of the state, i.e. the feature values (for binary approximators
    	 * this is abused to store the feature indices).
    	 */
    	protected T stateCoding;
    	
    	/**
    	 * Eligible options for the state to which these option values apply.
    	 */
    	protected List<? extends Option<?>> eligibleOpts;
    	
    	public FAOptionValues(State<? extends T> s) {
    		this.stateCoding = s.getData();
    		
    		eligibleOpts = s.getOptionSet();
    		if ( eligibleOpts == null )
    			eligibleOpts = optionPool;
    	}   	
    	    	
    	/* (non-Javadoc)
    	 * @see net.javlov.vf.OptionValues#getMaxOptions()
    	 */
    	@Override
    	public List<? extends Option<?>> getMaxOptions() {
    		double sum, maxValue = Double.NEGATIVE_INFINITY;
			List<Option<?>> maxOpts = new ArrayList<Option<?>>(eligibleOpts.size());
			for ( Option<?> o : eligibleOpts ) {
				sum = get(o);
				if ( sum > maxValue ) {
    				maxOpts.clear();
    				maxOpts.add(o);
    				maxValue = sum;
    			} else if ( sum == maxValue )
    				maxOpts.add(o);
			}
			return maxOpts;
    	}
    	
    	/* (non-Javadoc)
    	 * @see net.javlov.vf.OptionValues#getOptions()
    	 */
    	@Override
    	public List<? extends Option<?>> getOptions() {
    		return eligibleOpts;
    	}
    	
    	/* (non-Javadoc)
    	 * @see net.javlov.vf.OptionValues#getValues()
    	 */
    	@Override
    	public Collection<Value> getValues() {
    		List<Value> vals = new ArrayList<Value>(eligibleOpts.size());
    		for ( Option<?> o : eligibleOpts ) {
				vals.add( new Value(get(o)) );
    		}
    		return vals;
    	}
    	
    	/* (non-Javadoc)
    	 * @see net.javlov.vf.OptionValues#size()
    	 */
    	@Override
    	public int size() {
    		return eligibleOpts.size();
    	}
    	
    	/* (non-Javadoc)
    	 * @see net.javlov.vf.OptionValues#max()
    	 */
    	@Override
    	public double max() {
    		double sum, maxValue = Double.NEGATIVE_INFINITY;
			for ( Option<?> o : eligibleOpts ) {
				sum = get(o);
				if ( get(o) > maxValue )
    				maxValue = sum;
			}
			return maxValue;
    	}
    	
    	/**
    	 * Updates weights with non-zero traces. Called from the subclass.
    	 * 
    	 * @param o
    	 * @param delta
    	 */
    	protected void tracedWeightUpdate(Option<?> o, double delta) {
			double[][] tr = traces.get();
			int[] nonZero[] = traces.getLiveIndices(),
					numNonZero = traces.getNumNonZero();
			for ( int opt = 0; opt < optionPool.size(); opt++ ) {
				for ( int i = 0; i < numNonZero[opt]; i++ )
					weights[opt][nonZero[opt][i]] += delta*tr[opt][nonZero[opt][i]];
			}
    	}
    	
    	public String toString() {
    		return getValues().toString();
    	}

    }

	
	/**
	 * Linear Q-function approximator.
	 * @author Matthijs Snel
	 *
	 */
	public static class Linear extends QFA<double[]> {

		private static final long serialVersionUID = -2728185377552858032L;

		public Linear(int numFeats, List<? extends Option<?>> optionPool) {
			super(numFeats, optionPool);
		}

		@Override
		public OptionValues getValues(State<? extends double[]> s) {
			return new LinearOptionValues(s);
		}
		
		
		class LinearOptionValues extends FAOptionValues {

			private static final long serialVersionUID = 2448701162075647121L;

			public LinearOptionValues(State<? extends double[]> s) {
				super(s);
				stateCoding = Arrays.copyOf(stateCoding, stateCoding.length);
			}
			
	    	@Override
			public double get(Option<?> o) {	
				double sum = 0;
				for ( int i = 0; i < stateCoding.length; i++ )
					sum += weights[o.getID()][i]*stateCoding[i];
				return sum;
			}
	    	    
	    	@Override
			public void update(Option<?> o, double delta) {
	    		if ( traces != null ) {
	    			traces.update(stateCoding, o, this);
	    			tracedWeightUpdate(o,delta);
	    		}
	    		else {
	    			for ( int i = 0; i < stateCoding.length; i++ )
    					weights[o.getID()][i] += delta*stateCoding[i];
	    		}
			}
			
		}
	}
	
	
	
	
	public static class Binary extends QFA<int[]> {
		
		private static final long serialVersionUID = 687243198012436371L;

		public Binary(int numFeats, List<? extends Option<?>> optionPool) {
			super(numFeats, optionPool);
		}

		@Override
		public OptionValues getValues(State<? extends int[]> s) {
			return new BinaryOptionValues(s);
		}
		
		
		class BinaryOptionValues extends FAOptionValues {

			private static final long serialVersionUID = 5897499548638396860L;

			protected double[] featVals;
			
			public BinaryOptionValues(State<? extends int[]> s) {
				super(s);
				stateCoding = Arrays.copyOf(stateCoding, stateCoding.length);
				featVals = new double[stateCoding.length];
				for ( int i = 0; i < stateCoding.length; i++ )
					featVals[i] = 1; //binary approximator so values are always 1
			}
			
	    	@Override
			public double get(Option<?> o) {	
				double sum = 0;
				for ( int i = 0; i < stateCoding.length; i++ )
					sum += weights[o.getID()][stateCoding[i]];
				return sum;
			}
	    	    
	    	@Override
			public void update(Option<?> o, double delta) {
	    		if ( traces != null ) {
	    			traces.update(featVals, o, stateCoding, this);
	    			tracedWeightUpdate(o,delta/stateCoding.length);
	    		}
	    		else {
	    			for ( int i = 0; i < stateCoding.length; i++ )
    					weights[o.getID()][stateCoding[i]] += delta/stateCoding.length;
	    		}
			}
			
		}
	}
}
